{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {},
   "source": [
    "from rich.pretty import pprint\n",
    "\n",
    "from sagemaker.ai_registry.air_constants import REWARD_FUNCTION, REWARD_PROMPT\n",
    "from sagemaker.ai_registry.dataset import DataSet, CustomizationTechnique\n",
    "from sagemaker.ai_registry.evaluator import Evaluator"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "# Configure AWS credentials and region\n",
    "#! ada credentials update --provider=isengard --account=<> --role=Admin --profile=default --once\n",
    "#! aws configure set region us-west-2"
   ],
   "id": "665a0e71fef89bde"
  },
  {
   "cell_type": "markdown",
   "id": "28caba460ecd23c1",
   "metadata": {},
   "source": [
    "## DataSets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3066c60ef19a4e56",
   "metadata": {},
   "source": [
    "#### Create\n",
    "- DataSet input format depends on Customization technique\n",
    "- If no customization technique is provide, client side validation would be skipped\n",
    "- Provide a source (it can be local file path or S3 URL)"
   ]
  },
  {
   "cell_type": "code",
   "id": "2234f21780b91625",
   "metadata": {},
   "source": [
    "\n",
    "# 1. S3 Data source\n",
    "dataset = DataSet.create(\n",
    "            name=\"sdkv3-gen-ds2\",\n",
    "            source=\"s3://sdk-air-test-bucket/datasets/training-data/jamjee-sft-ds1.jsonl\",\n",
    "            customization_technique=CustomizationTechnique.SFT\n",
    "        )"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "61f55698ab27d70a",
   "metadata": {},
   "source": [
    "# 2. local dataset file source\n",
    "# ------------------------------------\n",
    "# To remove this line post testing/dogfooding : Sample source https://quip-amazon.com/hXbKA1U0aKTL/Model-Customisation-Bug-Bash#temp:s:temp:C:bYf1df6d6a2346e4fea8eb89d6c9;temp:C:bYf4ecae019198f4eb8940daf7f8\n",
    "# Download dataset from above link locally and provide data_location as local path.\n",
    "# Or, Upload the file to an accessible S3 location and provide S3 URI below as data_location.\n",
    "\n",
    "dataset = DataSet.create(\n",
    "            name=\"my-rlvr-ds1\",\n",
    "            source=\"/Volumes/workplace/sagemaker-python-sdk-staging/recipes-data/rlvr/train_256.jsonl\",\n",
    "            customization_technique=CustomizationTechnique.RLVR\n",
    "        )"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "ee2980471f8ae0c0",
   "metadata": {},
   "source": [
    "# Refreshes status from hub\n",
    "dataset.refresh()\n",
    "pprint(dataset.__dict__)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "30c1b17ad232110b",
   "metadata": {},
   "source": [
    "versions = dataset.get_versions()\n",
    "pprint(versions.__dict__)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "332be046d91fcefc",
   "metadata": {},
   "source": [
    "# delete specific version\n",
    "dataset.delete(version=\"0.0.4\")\n",
    "#dataset.delete(version=\"use a version from versions\")\n",
    "#pprint(versions)\n",
    "# specified deleted version should not be part of output"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "510d1a015e7a565c",
   "metadata": {},
   "source": [
    "# deletes all versions of this dataset by default\n",
    "dataset.delete()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "ca8f78c35ea9bf99",
   "metadata": {},
   "source": [
    "#### List DataSet"
   ]
  },
  {
   "cell_type": "code",
   "id": "d89a8741dd64f92e",
   "metadata": {},
   "source": [
    "#Optional max_results argument for pagination or else use default config\n",
    "datasets = DataSet.get_all(max_results=2)\n",
    "for dataset in datasets:\n",
    "    pprint(dataset)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "d8c16c305e1957bf",
   "metadata": {},
   "source": [
    "#### Use an existing DataSet"
   ]
  },
  {
   "cell_type": "code",
   "id": "572d4184cf42c7fa",
   "metadata": {},
   "source": [
    "# Use a dataset from iterator\n",
    "dataset = next(DataSet.get_all(max_results=2))\n",
    "for dataset in datasets:\n",
    "    pprint(dataset.__dict__)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "ae056f626cd7e931",
   "metadata": {},
   "source": [
    "# Use a dataset by name\n",
    "dataset = DataSet.get(name=\"sdkv3-gen-ds2\")\n",
    "pprint(dataset)\n",
    "\n",
    "# We can do CRUD operation on this DataSet\n",
    "# e.g. dataset.delete()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "44d7a8150b4b7846",
   "metadata": {},
   "source": [
    "#Create a new version of this dataset\n",
    "dataset.create_version(source=\"s3://sdk-air-test-bucket/datasets/test_ds\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "ba3ae7101c5281de",
   "metadata": {},
   "source": [
    "versions = dataset.get_versions()\n",
    "pprint(versions)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a73d88d38a2d5ba3",
   "metadata": {},
   "source": [
    "## Evaluator"
   ]
  },
  {
   "cell_type": "code",
   "id": "2d0ff33265d2c8dd",
   "metadata": {},
   "source": [
    "# Method : Lambda\n",
    "evaluator = Evaluator.create(\n",
    "    name = \"sdk-new-rf11\",\n",
    "    source=\"arn:aws:lambda:us-west-2:<>:function:sm-eval-vinayshm-rlvr-llama-321b-instruct-v1-<>8\",\n",
    "    type=REWARD_FUNCTION\n",
    "\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "ab2896e0b68b9384",
   "metadata": {},
   "source": [
    "# Method : BYOC\n",
    "\n",
    "evaluator = Evaluator.create(\n",
    "    name = \"eval-lambda-test\",\n",
    "    source=\"/Volumes/workplace/sagemaker-python-sdk-staging/recipes-data/eval_lambda_1.py\",\n",
    "    type = REWARD_FUNCTION\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "813243a997e3946b",
   "metadata": {},
   "source": [
    "# Reward Prompt\n",
    "# ------------------------------------\n",
    "# To remove this line post testing/dogfooding : Sample source https://quip-amazon.com/hXbKA1U0aKTL/Model-Customisation-Bug-Bash#temp:s:temp:C:bYf5c2e9e77efea4868b0420892a;temp:C:bYf4ecae019198f4eb8940daf7f8\n",
    "# Download prompt from above link locally and provide prompt_source as local path.\n",
    "# Or, Upload the file to a accessible S3 location and provide S3 URI below as prompt_source.\n",
    "\n",
    "evaluator = Evaluator.create(\n",
    "    name = \"jamj-rp2\",\n",
    "    source=\"/Users/jamjee/workplace/hubpuller/prompt/custom_prompt.jinja\",\n",
    "    type = REWARD_PROMPT\n",
    ")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "a7aef9b8a54766eb",
   "metadata": {},
   "source": [
    "# Optional wait, by default we have wait = True during create call.\n",
    "evaluator.wait()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "13ff6d34eab34a07",
   "metadata": {},
   "source": [
    "evaluator.refresh()\n",
    "pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "345214df-f320-4de0-ba97-860429f1f5bb",
   "metadata": {},
   "source": [
    "# Optional max_results for pagination\n",
    "evaluators = Evaluator.get_all(max_results=2)\n",
    "for evaluator in evaluators:\n",
    "    pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "b0f2cb26d5bb9a08",
   "metadata": {},
   "source": [
    "# Get evaluators by type\n",
    "evaluators = Evaluator.get_all(type='RewardPrompt', max_results=2)\n",
    "for evaluator in evaluators:\n",
    "    pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "1c62ec2f94eb9ac5",
   "metadata": {},
   "source": [
    "# Get an evaluator by name\n",
    "evaluator = Evaluator.get(name=\"sdk-new-rf11\")\n",
    "pprint(evaluator)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "b1a2154e870e623c",
   "metadata": {},
   "source": [
    "evaluator.create_version(source=evaluator.reference)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "72faf70127208509",
   "metadata": {},
   "source": [
    "versions = evaluator.get_versions()\n",
    "pprint(versions)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "0dc1107a-126b-4484-9639-07ba5de4ade6",
   "metadata": {},
   "source": [
    "# delete evaluator, option version argument or delete all versions.\n",
    "evaluator.delete()"
   ],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
